# starter助理 通过问一系列问题来引导用户创建自己的化身 并把用户的回答整理为化身可以接受的json格式 作为化身的提示词
# langchain

from langchain.agents import AgentExecutor
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI  # 从 langchain_openai 导入
from langchain.agents.format_scratchpad import format_to_openai_functions
from langchain.agents.output_parsers import OpenAIFunctionsAgentOutputParser
from langchain.tools import Tool, format_tool_to_openai_function  # 添加 format_tool_to_openai_function
from langchain.memory import ConversationBufferMemory  # 改用 ConversationBufferMemory
from pydantic.v1 import BaseModel  # 使用 v1 版本
from typing import List
from . import ai_base

# 修改工具函数
def upload_image(text: str = "") -> str:
    """引导用户上传外观图片"""
    return """请记住稍后在创建化身时需要上传一张您的照片作为外观。
照片要求：
1. 清晰的正面照
2. 光线充足
3. 建议使用证件照或形象照
4. 图片格式支持: jpg, png
现在让我们继续了解您的其他信息。"""

def upload_voice(text: str = "") -> str:
    """引导用户上传声音样本"""
    return """请记住稍后在创建化身时需要录制一段语音作为声音样本。
语音要求：
1. 清晰的朗读声音
2. 建议录制30秒左右
3. 可以读一段自我介绍
4. 音频格式支持: mp3, wav
现在让我们继续了解您的其他信息。"""

def create_avatar_profile(
    personality: str = "",
    interests: str = "",
    values: str = "",
    speaking_style: str = "",
    sample_dialogue: str = ""
) -> dict:
    """创建化身档案，整理用户的性格特征等信息
    Args:
        personality: 性格特征
        interests: 兴趣爱好
        values: 价值观
        speaking_style: 说话风格
        sample_dialogue: 示例对话
    Returns:
        dict: 化身档案
    """
    return {
        "personality": personality,
        "interests": interests,
        "values": values,
        "speaking_style": speaking_style,
        "sample_dialogue": sample_dialogue
    }

# 定义工具列表
# tools = [
#     Tool(
#         name="upload_image",
#         func=upload_image,
#         description="当需要用户上传外观图片时使用。可以提供提示文本。"
#     ),
#     Tool(
#         name="upload_voice", 
#         func=upload_voice,
#         description="当需要用户上传声音样本时使用。可以提供提示文本。"
#     ),
#     Tool(
#         name="create_avatar_profile",
#         func=create_avatar_profile,
#         description="当收集到用户的性格特征等信息时，用于创建化身档案。需要提供相应的信息字段。"
#     )
# ]

tools = []

llm = ai_base.llm

# 创建 memory - 使用简单的 ConversationBufferMemory
memory = ConversationBufferMemory(
    return_messages=True,
    memory_key='chat_history'
)

# 修改提示词模板
prompt = ChatPromptTemplate.from_messages([
    ("system", """你是一个友好的助手,负责通过对话了解用户并帮助创建他们的虚拟化身。

    你需要通过对话完成以下任务:
    1. 提示用户上传外观图片
    2. 提示用户上传声音样本
    3. 通过提问了解用户的性格特征、兴趣爱好、价值观等信息
    4. 将收集到的信息整理为化身档案

    请保持对话自然亲切，一步一步引导用户。不要一次性问太多问题。
    """),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{input}"),
    MessagesPlaceholder(variable_name="agent_scratchpad"),
])

# 创建 agent
llm_with_tools = llm.bind(
    functions=[format_tool_to_openai_function(t) for t in tools]
)

agent = (
    {
        "input": lambda x: x["input"],
        "chat_history": lambda x: memory.load_memory_variables({})["chat_history"],
        "agent_scratchpad": lambda x: format_to_openai_functions(x["intermediate_steps"])
    }
    | prompt
    | llm_with_tools
    | OpenAIFunctionsAgentOutputParser()
)

agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    memory=memory,
    max_iterations=1,  # 只执行一次
    early_stopping_method="force",  # 强制在第一次响应后停止
    verbose=True,
    handle_parsing_errors=True,
    return_intermediate_steps=False
)

# 修改chat_with_starter函数
async def chat_with_starter(message: str) -> str:
    """与starter助理对话"""
    try:
        response = agent_executor.invoke({
            "input": message
        })
        return response["output"] if isinstance(response, dict) else str(response)
    except Exception as e:
        print(f"Starter agent error: {str(e)}")
        return f"对话出现错误: {str(e)}"
